package tom.jma;

import javassist.*;
import tom.log.Output;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.lang.instrument.ClassFileTransformer;
import java.net.URL;
import java.security.ProtectionDomain;
import java.util.*;

/**
 * read perf method files in classpath: jma.properties <br>
 * <code>
 * jma.properties example: <br>
 * tom.demo.Hello=hi,say,ok <br>
 * tom.demo.World=world <br>
 *
 * </code>
 *
 * @author tom
 * @date 2021/12/15 09:52
 **/
public class MethodPerfTransformer implements ClassFileTransformer {

    private final static String VAR_CODE_1 = "\nStringBuilder argStr= new StringBuilder();\nfor (int i =0; i<$args.length; i++) {argStr.append($args[i].getClass().getName()).append(\"[\").append($args[i]).append(\"]\").append(\",\");}\nif (argStr.length() > 0) { argStr.deleteCharAt(argStr.length() - 1);}";
    private final static String TIME_CODE_1 = "\nlong start = System.currentTimeMillis();\n";
    private final static String TIME_CODE_2 = "\nlong end = System.currentTimeMillis();\n";
    private final static String OUTPUT_CODE_1 = "\nif ((end-start) > %d) {tom.log.Output.print(\"Method: %s(\" + argStr + \") time-cost: \" + (end - start) + \"ms\"); }";
    private final static String JMA_PROPERTIES_RESOURCE_LOCATION = "jma.properties";
    private final static String JMA_PROPERTY_COST_THRESHOLD = "jma.cost.threshold";
    private final Map<String, List<String>> clazzMethods = new HashMap<>();
    private int jmaCostThreshold = -1;

    /**
     * @param jmaPropertiesLocation if null use: {@link #JMA_PROPERTIES_RESOURCE_LOCATION}
     */
    public MethodPerfTransformer(String jmaPropertiesLocation) {
        jmaPropertiesLocation = Objects.isNull(jmaPropertiesLocation) ? JMA_PROPERTIES_RESOURCE_LOCATION : jmaPropertiesLocation;
        try {
            Properties localProperties = new Properties();
            ClassLoader cl = MethodPerfTransformer.class.getClassLoader();
            URL url = (cl != null ? cl.getResource(jmaPropertiesLocation) :
                    ClassLoader.getSystemResource(jmaPropertiesLocation));
            if (Objects.isNull(url) && new File(jmaPropertiesLocation).exists()) {
                url = new File(jmaPropertiesLocation).getAbsoluteFile().toURI().toURL();
            }
            if (url != null) {
                Output.print("Found '" + jmaPropertiesLocation + "' file in local classpath");
                try (InputStream is = url.openStream()) {
                    localProperties.load(is);
                }
                Enumeration<?> propertyNames = localProperties.propertyNames();
                while (propertyNames.hasMoreElements()) {
                    String name = String.valueOf(propertyNames.nextElement());
                    String methodsStr = localProperties.getProperty(name);
                    if (Objects.nonNull(methodsStr) && methodsStr.length() > 0) {
                        List<String> methodList = Arrays.asList(methodsStr.split(","));
                        clazzMethods.computeIfAbsent(name, t -> new ArrayList<>()).addAll(methodList);
                    }
                }
                if (clazzMethods.containsKey(JMA_PROPERTY_COST_THRESHOLD)) {
                    this.jmaCostThreshold = Integer.parseInt(clazzMethods.get(JMA_PROPERTY_COST_THRESHOLD).get(0));
                }
            } else {
                Output.print("Could not found '" + jmaPropertiesLocation + "' file from local classpath.");
            }
        } catch (IOException ex) {
            Output.print("Could not load '" + jmaPropertiesLocation + "' file from local classpath: " + ex);
        }
    }

    @Override
    public byte[] transform(ClassLoader paramClassLoader, String paramClassName, Class<?> classBeingRedefined,
                            ProtectionDomain paramProtectionDomain, byte[] classfileBuffer) {
        paramClassName = paramClassName.replace("/", ".");
        if (clazzMethods.containsKey(paramClassName)) {
            try {
                Output.print("Find transfer bytecode class: " + paramClassName);
                CtClass ctclass = getCtClass(paramClassLoader, paramClassName);
                for (String methodName : clazzMethods.get(paramClassName)) {
                    CtMethod ctmethod = ctclass.getDeclaredMethod(methodName);
                    String originMethod = methodName + "$origin";
                    ctmethod.setName(originMethod);
                    CtMethod newMethod = CtNewMethod.copy(ctmethod, methodName, ctclass, null);
                    String returnStatement = "";
                    if (!"void".equalsIgnoreCase(ctmethod.getReturnType().getName())) {
                        returnStatement = "return ";
                    }
                    String methodBodyStr = "{" +
                            VAR_CODE_1 +
                            TIME_CODE_1 +
                            "try {" +
                            returnStatement + originMethod + "($$);\n" +
                            "} finally {" +
                            TIME_CODE_2 +
                            String.format(OUTPUT_CODE_1, this.jmaCostThreshold, methodName) +
                            "}" +
                            "}";
                    newMethod.setBody(methodBodyStr);
                    ctclass.addMethod(newMethod);
                }
                return ctclass.toBytecode();
            } catch (Exception e) {
                Output.print("Class bytecode transform failed: " + e);
                e.printStackTrace();
            }
        }
        return null;
    }

    private CtClass getCtClass(ClassLoader paramClassLoader, String paramClassName) throws NotFoundException {
        try {
            return ClassPool.getDefault().get(paramClassName);
        } catch (NotFoundException e) {
            Output.print("Class: '" + paramClassName + "' not found in default search path.");
            ClassPool.getDefault().appendClassPath(new LoaderClassPath(paramClassLoader));
        }
        return ClassPool.getDefault().get(paramClassName);
    }

}